import { FC, forwardRef, ForwardRefRenderFunction, useImperativeHandle, useState } from 'react'
import { View, Image } from '@tarojs/components'
import Checkbox from './checkbox'
import classNames from 'classnames'
import styles from './index.modules.less'
import expand_icon from './images/expend.png'
import all_block from './images/all_block.png'
import normal_icon from './images/normal.png'
import normal_acitve from './images/normal_acitve.png'
import { flatCheckedTreeNode, getSelectionStatus } from './const'
// 定义树节点类型
export interface TreeNode {
  id: string
  label: string
  children?: TreeNode[]
}

interface IProps {
  data: TreeNode[]
  onChange?: (flatCheckedTreeNode: TreeNode[]) => void
}

export interface TreeListInstanceRef {
  reset: () => void
}

// TreeList组件
const TreeList: ForwardRefRenderFunction<unknown, IProps> = ({ data, onChange }, ref) => {
  // 管理选中状态
  const [checkedMap, setCheckedMap] = useState<Record<string, boolean>>({})
  // 管理展开状态
  const [expandedMap, setExpandedMap] = useState<Record<string, boolean>>({})

  useImperativeHandle(ref, () => ({
    reset: () => {
      setCheckedMap({})
      setExpandedMap({})
      onChange?.([])
    }
  }))

  // 更新选中状态（支持父子联动）
  const handleCheckChange = (id: string, checked: boolean) => {
    const updateChecked = (nodes: TreeNode[]) => {
      return nodes.reduce((map, node) => {
        map[node.id] = checked
        if (node.children) {
          Object.assign(map, updateChecked(node.children))
        }
        return map
      }, {} as Record<string, boolean>)
    }

    const newState = {
      ...checkedMap,
      ...updateChecked(findNodeAndChildren(data, id))
    }

    setCheckedMap(newState)
    const flatTreeNode = flatCheckedTreeNode(newState, data)
    onChange?.(flatTreeNode)
  }

  // 查找当前节点及其所有子节点
  const findNodeAndChildren = (nodes: TreeNode[], id: string): TreeNode[] => {
    for (const node of nodes) {
      if (node.id === id) {
        return [node, ...(node.children || [])]
      }
      if (node.children) {
        const result = findNodeAndChildren(node.children, id)
        if (result.length) return result
      }
    }
    return []
  }

  // 切换节点的展开状态
  const toggleExpand = (id: string) => {
    setExpandedMap((prev) => ({
      ...prev,
      [id]: !prev[id]
    }))
  }

  // 全选或取消所有节点
  const handleSelectAll = () => {
    const selectionStatus = getSelectionStatus(data, checkedMap)

    const updateChecked = (nodes: TreeNode[], checked: boolean) => {
      return nodes.reduce((map, node) => {
        map[node.id] = checked
        if (node.children) {
          Object.assign(map, updateChecked(node.children, checked))
        }
        return map
      }, {} as Record<string, boolean>)
    }

    // 根据当前状态决定是全选还是取消
    const newCheckedMap = selectionStatus !== 'all' ? updateChecked(data, true) : {}
    setCheckedMap(newCheckedMap)
    const flatTreeNode = flatCheckedTreeNode(newCheckedMap, data)
    onChange?.(flatTreeNode)
  }

  const renderToggleExpand = (id: string) => {
    const isExpanded = !!expandedMap[id]
    return (
      <Image
        src={expand_icon}
        className={classNames(styles.expandedIcon, !isExpanded && styles.expanded)}
        onClick={(event) => {
          event.stopPropagation()
          toggleExpand(id)
        }}
      />
    )
  }

  // 渲染树节点
  const renderTree = (nodes: TreeNode[], level = 0) => {
    return nodes.map((node) => {
      const isExpanded = !!expandedMap[node.id]
      const hasChildren = Array.isArray(node.children) && node.children.length > 0

      return (
        <View key={node.id} className={classNames(styles.treeNode, styles[`level-${level}`])}>
          <View style={{ marginBottom: 20 }}>
            {/* 复选框 */}
            <Checkbox value={!!checkedMap[node.id]} onChange={(newValue) => handleCheckChange(node.id, newValue)}>
              {/* 节点标签 */}
              <span>{node.label}</span>
              {/* 展开/收起按钮 */}
              {hasChildren && renderToggleExpand(node.id)}
            </Checkbox>
          </View>

          {/* 子节点 */}
          {hasChildren && isExpanded && renderTree(node.children!, level + 1)}
        </View>
      )
    })
  }

  const renderAllIcon = () => {
    const selectionStatus = getSelectionStatus(data, checkedMap)
    const icon = selectionStatus === 'all' ? normal_acitve : selectionStatus === 'partial' ? all_block : normal_icon
    return <Image src={icon} className={styles.allIcon} />
  }

  return (
    <View className={styles.treeList}>
      {renderTree(data)}
      <View className={styles.itemCenter} onClick={handleSelectAll}>
        {renderAllIcon()}
        全部
      </View>
    </View>
  )
}

export default forwardRef(TreeList)
